#!/usr/bin/env python3
"""
安全测试客户端
测试增强的安全功能
"""

import requests
import json
import time
import threading
from datetime import datetime
from utils import EncryptionUtil

class SecurityTestClient:
    """安全测试客户端"""
    
    def __init__(self, base_url='http://localhost:8000'):
        self.base_url = base_url
        self.session = requests.Session()
        self.encryption_util = EncryptionUtil('your-encryption-key-here')
        self.test_results = []
        
    def log_result(self, test_name, success, message, details=None):
        """记录测试结果"""
        result = {
            'test_name': test_name,
            'success': success,
            'message': message,
            'details': details,
            'timestamp': datetime.now().isoformat()
        }
        self.test_results.append(result)
        status = "✅" if success else "❌"
        print(f"{status} {test_name}: {message}")
        
    def test_rate_limiting(self):
        """测试频率限制"""
        print("\n🔄 测试频率限制...")
        
        # 测试登录频率限制
        login_data = {
            'username': 'test_user',
            'password': 'wrong_password'
        }
        
        success_count = 0
        blocked_count = 0
        
        for i in range(10):
            try:
                response = self.session.post(
                    f'{self.base_url}/api/auth/login',
                    json=login_data,
                    timeout=5
                )
                
                if response.status_code == 429:
                    blocked_count += 1
                elif response.status_code in [200, 401]:
                    success_count += 1
                    
            except Exception as e:
                print(f"请求失败: {str(e)}")
                
            time.sleep(0.5)
        
        if blocked_count > 0:
            self.log_result(
                "频率限制测试",
                True,
                f"成功触发频率限制，{success_count}次成功，{blocked_count}次被阻止"
            )
        else:
            self.log_result(
                "频率限制测试",
                False,
                "未能触发频率限制"
            )
    
    def test_input_validation(self):
        """测试输入验证"""
        print("\n🛡️ 测试输入验证...")
        
        # 测试SQL注入
        sql_injection_payloads = [
            "'; DROP TABLE users; --",
            "' OR '1'='1",
            "admin'--",
            "' UNION SELECT * FROM users --"
        ]
        
        blocked_sql = 0
        for payload in sql_injection_payloads:
            try:
                response = self.session.post(
                    f'{self.base_url}/api/auth/login',
                    json={
                        'username': payload,
                        'password': 'test'
                    },
                    timeout=5
                )
                
                if response.status_code == 400:
                    blocked_sql += 1
                    
            except Exception:
                pass
        
        # 测试XSS
        xss_payloads = [
            "<script>alert('xss')</script>",
            "javascript:alert('xss')",
            "<img src=x onerror=alert('xss')>",
            "<svg onload=alert('xss')>"
        ]
        
        blocked_xss = 0
        for payload in xss_payloads:
            try:
                response = self.session.post(
                    f'{self.base_url}/api/auth/send-sms',
                    json={
                        'phone': payload
                    },
                    timeout=5
                )
                
                if response.status_code == 400:
                    blocked_xss += 1
                    
            except Exception:
                pass
        
        total_blocked = blocked_sql + blocked_xss
        total_tests = len(sql_injection_payloads) + len(xss_payloads)
        
        if total_blocked > total_tests * 0.5:  # 至少阻止50%的攻击
            self.log_result(
                "输入验证测试",
                True,
                f"成功阻止 {total_blocked}/{total_tests} 个恶意输入"
            )
        else:
            self.log_result(
                "输入验证测试",
                False,
                f"仅阻止 {total_blocked}/{total_tests} 个恶意输入"
            )
    
    def test_encryption(self):
        """测试加密功能"""
        print("\n🔐 测试加密功能...")
        
        try:
            # 测试加密请求
            test_data = {
                'username': 'test_user',
                'password': 'test_password'
            }
            
            encrypted_data = self.encryption_util.encrypt_json(test_data)
            
            response = self.session.post(
                f'{self.base_url}/api/auth/login',
                json={
                    'encrypted': True,
                    'data': encrypted_data
                },
                timeout=5
            )
            
            if response.status_code in [200, 401]:
                response_data = response.json()
                if response_data.get('encrypted'):
                    # 尝试解密响应
                    decrypted = self.encryption_util.decrypt_json(response_data['data'])
                    self.log_result(
                        "加密功能测试",
                        True,
                        "加密请求和响应正常工作"
                    )
                else:
                    self.log_result(
                        "加密功能测试",
                        False,
                        "响应未加密"
                    )
            else:
                self.log_result(
                    "加密功能测试",
                    False,
                    f"加密请求失败: {response.status_code}"
                )
                
        except Exception as e:
            self.log_result(
                "加密功能测试",
                False,
                f"加密测试异常: {str(e)}"
            )
    
    def test_concurrent_requests(self):
        """测试并发请求处理"""
        print("\n⚡ 测试并发请求...")
        
        def make_request():
            try:
                response = self.session.get(f'{self.base_url}/health', timeout=5)
                return response.status_code == 200
            except:
                return False
        
        # 创建多个线程同时发送请求
        threads = []
        results = []
        
        for i in range(20):
            thread = threading.Thread(target=lambda: results.append(make_request()))
            threads.append(thread)
            thread.start()
        
        # 等待所有线程完成
        for thread in threads:
            thread.join()
        
        success_rate = sum(results) / len(results) if results else 0
        
        if success_rate > 0.8:  # 80%成功率
            self.log_result(
                "并发请求测试",
                True,
                f"并发处理成功率: {success_rate:.2%}"
            )
        else:
            self.log_result(
                "并发请求测试",
                False,
                f"并发处理成功率过低: {success_rate:.2%}"
            )
    
    def test_security_headers(self):
        """测试安全响应头"""
        print("\n🛡️ 测试安全响应头...")
        
        try:
            response = self.session.get(f'{self.base_url}/health', timeout=5)
            
            required_headers = [
                'X-Content-Type-Options',
                'X-Frame-Options',
                'X-XSS-Protection',
                'Strict-Transport-Security',
                'Content-Security-Policy'
            ]
            
            present_headers = []
            for header in required_headers:
                if header in response.headers:
                    present_headers.append(header)
            
            if len(present_headers) >= len(required_headers) * 0.8:  # 至少80%的安全头
                self.log_result(
                    "安全响应头测试",
                    True,
                    f"存在 {len(present_headers)}/{len(required_headers)} 个安全头"
                )
            else:
                self.log_result(
                    "安全响应头测试",
                    False,
                    f"缺少安全头，仅有 {len(present_headers)}/{len(required_headers)} 个"
                )
                
        except Exception as e:
            self.log_result(
                "安全响应头测试",
                False,
                f"测试异常: {str(e)}"
            )
    
    def test_api_endpoints(self):
        """测试API端点可用性"""
        print("\n🔗 测试API端点...")
        
        endpoints = [
            ('/health', 'GET'),
            ('/', 'GET'),
            ('/api/auth/send-sms', 'POST'),
            ('/api/security/monitor/status', 'GET')  # 需要认证，应该返回401
        ]
        
        working_endpoints = 0
        
        for endpoint, method in endpoints:
            try:
                if method == 'GET':
                    response = self.session.get(f'{self.base_url}{endpoint}', timeout=5)
                else:
                    response = self.session.post(f'{self.base_url}{endpoint}', json={}, timeout=5)
                
                # 200, 400, 401都算正常响应
                if response.status_code in [200, 400, 401]:
                    working_endpoints += 1
                    
            except Exception:
                pass
        
        if working_endpoints >= len(endpoints) * 0.8:
            self.log_result(
                "API端点测试",
                True,
                f"{working_endpoints}/{len(endpoints)} 个端点正常响应"
            )
        else:
            self.log_result(
                "API端点测试",
                False,
                f"仅 {working_endpoints}/{len(endpoints)} 个端点正常响应"
            )
    
    def run_all_tests(self):
        """运行所有安全测试"""
        print("🚀 开始安全测试...")
        print("=" * 50)
        
        # 检查服务器是否可用
        try:
            response = self.session.get(f'{self.base_url}/health', timeout=5)
            if response.status_code != 200:
                print("❌ 服务器不可用，测试终止")
                return
        except Exception as e:
            print(f"❌ 无法连接到服务器: {str(e)}")
            return
        
        print("✅ 服务器连接正常")
        
        # 运行各项测试
        self.test_api_endpoints()
        self.test_security_headers()
        self.test_encryption()
        self.test_input_validation()
        self.test_rate_limiting()
        self.test_concurrent_requests()
        
        # 生成测试报告
        self.generate_report()
    
    def generate_report(self):
        """生成测试报告"""
        print("\n" + "=" * 50)
        print("📊 安全测试报告")
        print("=" * 50)
        
        total_tests = len(self.test_results)
        passed_tests = sum(1 for result in self.test_results if result['success'])
        
        print(f"总测试数: {total_tests}")
        print(f"通过测试: {passed_tests}")
        print(f"失败测试: {total_tests - passed_tests}")
        print(f"通过率: {passed_tests/total_tests:.2%}")
        
        print("\n详细结果:")
        for result in self.test_results:
            status = "✅" if result['success'] else "❌"
            print(f"{status} {result['test_name']}: {result['message']}")
        
        # 保存报告到文件
        report_data = {
            'timestamp': datetime.now().isoformat(),
            'summary': {
                'total_tests': total_tests,
                'passed_tests': passed_tests,
                'failed_tests': total_tests - passed_tests,
                'pass_rate': passed_tests/total_tests
            },
            'results': self.test_results
        }
        
        with open('security_test_report.json', 'w', encoding='utf-8') as f:
            json.dump(report_data, f, ensure_ascii=False, indent=2)
        
        print(f"\n📄 详细报告已保存到: security_test_report.json")

def main():
    """主函数"""
    client = SecurityTestClient()
    client.run_all_tests()

if __name__ == '__main__':
    main()